package qq2564874169.tuaotuao.fx.orm.jdbc;



import qq2564874169.tuaotuao.fx.FxAssert;

import javax.sql.DataSource;
import java.io.PrintWriter;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.SQLFeatureNotSupportedException;
import java.util.HashMap;
import java.util.Map;
import java.util.logging.Logger;

public class MultiDataSource implements DataSource, AutoCloseable {
    private Map<String, DataSource> dss = new HashMap<>();
    private ThreadLocal<DataSource> ds = new ThreadLocal<>();
    private String defaultKey;

    public void add(String key, DataSource dataSource) {
        dss.put(key, dataSource);
    }

    public void set(String key) {
        FxAssert.isTrue(dss.containsKey(key), "不存在key为" + key + "的数据源");
        ds.set(dss.get(key));
    }

    public void reset() {
        ds.remove();
    }

    public void setDefaultKey(String key) {
        defaultKey = key;
    }

    public DataSource get() {
        DataSource dataSource = ds.get();
        if (dataSource == null) {
            dataSource = dss.get(defaultKey);
        }
        return dataSource;
    }

    @Override
    public Connection getConnection() throws SQLException {
        return get().getConnection();
    }

    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        return get().getConnection(username, password);
    }

    @Override
    public <T> T unwrap(Class<T> iface) throws SQLException {
        return get().unwrap(iface);
    }

    @Override
    public boolean isWrapperFor(Class<?> iface) throws SQLException {
        return get().isWrapperFor(iface);
    }

    @Override
    public PrintWriter getLogWriter() throws SQLException {
        return get().getLogWriter();
    }

    @Override
    public void setLogWriter(PrintWriter out) throws SQLException {
        get().setLogWriter(out);
    }

    @Override
    public void setLoginTimeout(int seconds) throws SQLException {
        get().setLoginTimeout(seconds);
    }

    @Override
    public int getLoginTimeout() throws SQLException {
        return get().getLoginTimeout();
    }

    @Override
    public Logger getParentLogger() throws SQLFeatureNotSupportedException {
        return get().getParentLogger();
    }

    @Override
    public void close() {
        if (ds != null) {
            ds.remove();
            ds = null;
        }
        if (dss != null) {
            for (Object obj : dss.values()) {
                if (obj instanceof AutoCloseable) {
                    try {
                        ((AutoCloseable) obj).close();
                    } catch (Exception ignored) {

                    }
                }
            }
            dss.clear();
            dss = null;
        }
        defaultKey = null;
    }
}
